from projects.mmdet3d_plugin.ops import DeformableAggregationFunction
from projects.mmdet3d_plugin.ops import deformable_aggregation_function
import torch
import torch.nn as nn
import torchvision
import torch.onnx
import onnxruntime
from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #这里创建一个层对DeformableAggregationFunction 实例化是不行的，
        # 对于自定义的function直接在forward 调用就行
    def forward(self, mc_ms_feat,spatial_shape,scale_start_index,sampling_location,weights):
        x = deformable_aggregation_function(mc_ms_feat,spatial_shape,scale_start_index,sampling_location,weights)
        return x


def validate_onnx():
    input = torch.rand(1, 3, 5, 5)

    # PyTorch的推理
    model = Model()
    x     = model(input)
    print("result from Pytorch is :", x)

    # onnxruntime的推理
    # sess  = onnxruntime.InferenceSession('../models/sample-deformable-conv.onnx')
    # x     = sess.run(None, {'input0': input.numpy()})
    # print("result from onnx is:    ", x)

def infer():
    # import pdb;pdb.set_trace()
    mc_ms_feat = torch.load("input/mc_ms_feat.pt")
    mc_ms_feat.requires_grad_(False)
    # mc_ms_feat = torch.rand(1, 89760, 256)

    spatial_shape = torch.load("input/spatial_shape.pt")
    spatial_shape.requires_grad_(False)
    # spatial_shape=torch.rand(6, 4, 2)

    scale_start_index = torch.load("input/scale_start_index.pt")
    scale_start_index.requires_grad_(False)
    # scale_start_index=torch.rand(6,4)

    sampling_location = torch.load("input/sampling_location.pt")
    sampling_location.requires_grad_(False)
    # sampling_location=torch.rand(1, 900, 13, 6, 2)

    weights = torch.load("input/weights.pt")
    weights.requires_grad_(False)
    # weights=torch.rand(1, 900, 13, 6, 4, 8)

    
    model = Model()
    model.cuda()
    x = model(mc_ms_feat,spatial_shape,scale_start_index,sampling_location,weights)
    # print("input is: ", input.data)
    print("result is: ", x.data)

    model.eval()

    file    = "../models/sample-daf.onnx"
    torch.onnx.export(
        model         = model, 
        args          = (mc_ms_feat,spatial_shape,scale_start_index,sampling_location,weights,),
        f             = file,
        input_names   = ["mc_ms_feat","spatial_shape","scale_start_index","sampling_location","weights"],
        output_names  = ["output0"],
        opset_version = 16)
    print("Finished normal onnx export")

if __name__ == "__main__":
    infer()
    # export_norm_onnx()
    # validate_onnx()
